use std::net::SocketAddr;

use axum::{
    extract::{
        ws::{Message, WebSocket},
        ConnectInfo, WebSocketUpgrade,
    },
    response::IntoResponse,
    routing::get,
    Router,
};
use futures::{SinkExt, StreamExt};
use tracing::info;

#[tokio::main]
async fn main() {
    std::env::set_var("RUST_LOG", "axum_ws=debug");
    tracing_subscriber::fmt::init();

    let app = Router::new().route("/ws", get(ws_handler));

    let lst = tokio::net::TcpListener::bind("0.0.0.0:56789")
        .await
        .unwrap();
    info!("listeting on {}", lst.local_addr().unwrap());
    axum::serve(lst, app.into_make_service_with_connect_info::<SocketAddr>())
        .await
        .unwrap();
}

async fn ws_handler(
    ws: WebSocketUpgrade,
    ConnectInfo(addr): ConnectInfo<SocketAddr>,
) -> impl IntoResponse {
    ws.on_upgrade(move |socket| handle_socket(socket, addr))
}

async fn handle_socket(socket: WebSocket, who: SocketAddr) {
    let (mut sender, mut receiver) = socket.split();

    while let Some(Ok(msg)) = receiver.next().await {
        if let Message::Close(_) = msg {
            info!("closed from client {who}");
            break;
        }
        match sender.send(msg).await {
            Ok(_) => info!("send msg ok {who}"),
            Err(err) => {
                info!("failed to send msg: {:?} {who}", err);
                break;
            }
        };
    }
}
